#=
一些常用的功能
=#

"""
一维链的哈密顿量
"""
function lattice_ionic_chain(::Type{T}, L, hop_t::T, delta::T; periodic=true) where T
    hk = zeros(T, L, L)
    for i in Base.OneTo(L-1)
        hk[i, i+1] = hop_t
        hk[i+1, i] = adjoint(hop_t)
    end
    if periodic && L > 2
        hk[1, L] = hop_t
        hk[L, 1] = adjoint(hop_t)
    end
    #
    for i in Base.OneTo(L)
        hk[i, i] += (-1)^(i+1)*delta
    end
    return hk
end


lattice_chain(L, hop_t; periodic=true) = lattice_ionic_chain(ComplexF64, L, complex(hop_t), 0.0+0.0im; periodic=periodic)


"""
六角晶格的哈密顿量
"""
function lattice_hexagonal(::Type{T}, L, hop_t::T; periodic=true, λ::T=T(0.0)) where T
    if !periodic
        throw(error("not impl"))
    end
    #
    Nx = 2*L^2
    htmp = zeros(T, Nx, Nx)
    ucmap = zeros(Int, L, L)
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        ucmap[ux, uy] = L*(uy-1) + ux
    end; end
    #
    println(ucmap)
    #
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        Aidx = 2*ucmap[ux, uy] - 1
        #自己子格子的
        Bidx = 2*ucmap[ux, uy]
        htmp[Aidx, Bidx] = hop_t
        htmp[Bidx, Aidx] = adjoint(hop_t)
        #-y
        #println(uy, @↺ L uy-1)
        Bidx = 2*ucmap[ux, @↺ L uy-1]
        htmp[Aidx, Bidx] = hop_t
        htmp[Bidx, Aidx] = adjoint(hop_t)
        #-y+x
        #println(uy, @↻(L,ux+1), @↺(L, @↻(L, ux+1)))
        Bidx = 2*ucmap[@↻(L, ux+1), @↺(L, uy-1)]
        htmp[Aidx, Bidx] = hop_t
        htmp[Bidx, Aidx] = adjoint(hop_t)
        #
    end; end
    #up和down
    hk = kron([1 0; 0 1], htmp)
    #
    htmp = zeros(T, Nx, Nx)
    σ_ud = [1, -im, 0]
    #up -> down
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        Aidx = 2*ucmap[ux, uy] - 1
        #自己格子的
        Bidx = 2*ucmap[ux, uy]
        dij = [0.0, sqrt(3)/3, 0.0]
        lv = σ_ud × dij
        htmp[Aidx, Bidx] += im*λ*lv[3]
        htmp[Bidx, Aidx] += -im*λ*lv[3]
        #-y
        Bidx = 2*ucmap[ux, @↺ L uy-1]
        dij = [-1/2, -sqrt(3)/6, 0.0]
        lv = σ_ud × dij
        htmp[Aidx, Bidx] += im*λ*lv[3]
        htmp[Bidx, Aidx] += -im*λ*lv[3]
        #-y+x
        Bidx = 2*ucmap[@↻(L, ux+1), @↺(L, uy-1)]
        dij = [1/2, -sqrt(3)/6, 0.0]
        lv = σ_ud × dij
        htmp[Aidx, Bidx] += im*λ*lv[3]
        htmp[Bidx, Aidx] += -im*λ*lv[3]
    end; end
    hk += kron([0 1; 0 0], htmp)
    #
    htmp = zeros(T, Nx, Nx)
    σ_du = [1, +im, 0]
    #dowm -> up
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        Aidx = 2*ucmap[ux, uy] - 1
        #自己格子的
        Bidx = 2*ucmap[ux, uy]
        dij = [0.0, sqrt(3)/3, 0.0]
        lv = σ_du × dij
        htmp[Aidx, Bidx] += im*λ*lv[3]
        htmp[Bidx, Aidx] += -im*λ*lv[3]
        #-y
        Bidx = 2*ucmap[ux, @↺ L uy-1]
        dij = [-1/2, -sqrt(3)/6, 0.0]
        lv = σ_du × dij
        htmp[Aidx, Bidx] += im*λ*lv[3]
        htmp[Bidx, Aidx] += -im*λ*lv[3]
        #-y+x
        Bidx = 2*ucmap[@↻(L, ux+1), @↺(L, uy-1)]
        dij = [1/2, -sqrt(3)/6, 0.0]
        lv = σ_du × dij
        htmp[Aidx, Bidx] += im*λ*lv[3]
        htmp[Bidx, Aidx] += -im*λ*lv[3]
    end; end
    hk += kron([0 0; 1 0], htmp)
    #
    @assert ishermitian(hk)
    return hk
end


function lattice_kagome(::Type{T}, L, hop_t::T; periodic=true, λR::T=T(0.0), λI::T=T(0.0)) where T
    if !periodic
        throw(error("not impl"))
    end
    #
    Nx = 3*L^2
    htmp = zeros(T, Nx, Nx)
    ucmap = zeros(Int, L, L)
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        ucmap[ux, uy] = L*(uy-1) + ux
    end; end
    #1811.08182
    println(ucmap)
    #
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        Aidx = 3*ucmap[ux, uy] - 2
        #A找2个C
        #自己元胞的
        Cidx = 3*ucmap[ux, uy]
        htmp[Aidx, Cidx] = hop_t
        htmp[Cidx, Aidx] = adjoint(hop_t)
        #-x-y的C
        Cidx = 3*ucmap[@↺(L, ux-1), @↺(L, uy-1)]
        htmp[Aidx, Cidx] = hop_t
        htmp[Cidx, Aidx] = adjoint(hop_t)
        #B找两个A
        Bidx = 3*ucmap[ux, uy] - 1
        #自己的A
        Aidx = 3*ucmap[ux, uy] - 2
        htmp[Bidx, Aidx] = hop_t
        htmp[Aidx, Bidx] = adjoint(hop_t)
        #+x的A
        Aidx = 3*ucmap[@↻(L, ux+1), uy] - 2
        htmp[Bidx, Aidx] = hop_t
        htmp[Aidx, Bidx] = adjoint(hop_t)
        #C找2个B
        Cidx = 3*ucmap[ux, uy]
        #自己的B
        Bidx = 3*ucmap[ux, uy] - 1
        htmp[Cidx, Bidx] = hop_t
        htmp[Bidx, Cidx] = adjoint(hop_t)
        #+y的B
        Bidx = 3*ucmap[ux, @↻(L, uy+1)] - 1
        htmp[Cidx, Bidx] = hop_t
        htmp[Bidx, Cidx] = adjoint(hop_t)
    end; end
    if isapprox(λR, 0.0) && isapprox(λI, 0.0)
        @assert ishermitian(htmp)
        return htmp
    end
    d1 = √3 / 2 * [0 1; 1 0] - 1 / 2 * [0 -im; im 0]
    d2 = [0 -im; im 0]
    d3 = -√3 / 2 * [0 1; 1 0] - 1 / 2 * [0 -im; im 0]
    σz = [1 0; 0 -1]
    #SOC
    hk = kron([1 0; 0 1], htmp)
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        Aidx = 3*ucmap[ux, uy] - 2
        #A找2个C
        #自己元胞的
        Cidx = 3*ucmap[ux, uy]
        Aup, Adn = Aidx, Aidx+Nx
        Cup, Cdn = Cidx, Cidx+Nx
        hk[[Aup, Adn], [Cup, Cdn]] += im*λR*d3
        hk[[Cup, Cdn], [Aup, Adn]] += adjoint(im*λR*d3)
        hk[[Aup, Adn], [Cup, Cdn]] += im*λI*σz
        hk[[Cup, Cdn], [Aup, Adn]] += adjoint(im*λI*σz)
        #-x-y的C
        Cidx = 3*ucmap[@↺(L, ux-1), @↺(L, uy-1)]
        Aup, Adn = Aidx, Aidx+Nx
        Cup, Cdn = Cidx, Cidx+Nx
        hk[[Aup, Adn], [Cup, Cdn]] += im*λR*d3
        hk[[Cup, Cdn], [Aup, Adn]] += adjoint(im*λR*d3)
        hk[[Aup, Adn], [Cup, Cdn]] += im*λI*σz
        hk[[Cup, Cdn], [Aup, Adn]] += adjoint(im*λI*σz)
        #B找两个A
        Bidx = 3*ucmap[ux, uy] - 1
        #自己的A
        Aidx = 3*ucmap[ux, uy] - 2
        Bup, Bdn = Bidx, Bidx+Nx
        Aup, Adn = Aidx, Aidx+Nx
        hk[[Bup, Bdn], [Aup, Adn]] += im*λR*d1
        hk[[Aup, Adn], [Bup, Bdn]] += adjoint(im*λR*d1)
        hk[[Bup, Bdn], [Aup, Adn]] += im*λI*σz
        hk[[Aup, Adn], [Bup, Bdn]] += adjoint(im*λI*σz)
        #+x的A
        Aidx = 3*ucmap[@↻(L, ux+1), uy] - 2
        Bup, Bdn = Bidx, Bidx+Nx
        Aup, Adn = Aidx, Aidx+Nx
        hk[[Bup, Bdn], [Aup, Adn]] += im*λR*d1
        hk[[Aup, Adn], [Bup, Bdn]] += adjoint(im*λR*d1)
        hk[[Bup, Bdn], [Aup, Adn]] += im*λI*σz
        hk[[Aup, Adn], [Bup, Bdn]] += adjoint(im*λI*σz)
        #C找2个B
        Cidx = 3*ucmap[ux, uy]
        #自己的B
        Bidx = 3*ucmap[ux, uy] - 1
        Cup, Cdn = Cidx, Cidx+Nx
        Bup, Bdn = Bidx, Bidx+Nx
        hk[[Cup, Cdn], [Bup, Bdn]] += im*λR*d2
        hk[[Bup, Bdn], [Cup, Cdn]] += adjoint(im*λR*d2)
        hk[[Cup, Cdn], [Bup, Bdn]] += im*λI*σz
        hk[[Bup, Bdn], [Cup, Cdn]] += adjoint(im*λI*σz)
        #+y的B
        Bidx = 3*ucmap[ux, @↻(L, uy+1)] - 1
        Cup, Cdn = Cidx, Cidx+Nx
        Bup, Bdn = Bidx, Bidx+Nx
        hk[[Cup, Cdn], [Bup, Bdn]] += im*λR*d2
        hk[[Bup, Bdn], [Cup, Cdn]] += adjoint(im*λR*d2)
        hk[[Cup, Cdn], [Bup, Bdn]] += im*λI*σz
        hk[[Bup, Bdn], [Cup, Cdn]] += adjoint(im*λI*σz)
    end; end
    @assert ishermitian(hk)
    return hk
end


"""
正方附加stagger
"""
function lattice_ionic_square(::Type{T}, L, hop_t::T, delta::T; periodic=true) where T
    Nx = L^2
    hk = zeros(T, Nx, Nx)
    ucmap = zeros(Int, L, L)
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        ucmap[ux, uy] = L*(uy-1) + ux
    end; end
    #
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        Aidx = ucmap[ux, uy]
        Bidx = ucmap[@↻(L, ux+1), uy]
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #
        Bidx = ucmap[ux, @↻(L, uy+1)]
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #
        if mod(ux+uy, 2) == 0
            hk[Aidx, Aidx] += delta
        else
            hk[Aidx, Aidx] += -delta
        end
    end; end
    return hk
end


"""
正方附加t'
"""
function lattice_tprim_square(::Type{T}, L, hop_t::T, tprim::T; periodic=true) where T
    Nx = L^2
    hk = zeros(T, Nx, Nx)
    ucmap = zeros(Int, L, L)
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        ucmap[ux, uy] = L*(uy-1) + ux
    end; end
    #
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        Aidx = ucmap[ux, uy]
        Bidx = ucmap[@↻(L, ux+1), uy]
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #
        Bidx = ucmap[ux, @↻(L, uy+1)]
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #
        Bidx = ucmap[@↻(L, ux+1), @↻(L, uy+1)]
        hk[Aidx, Bidx] += tprim
        hk[Bidx, Aidx] += adjoint(tprim)
        #
        Bidx = ucmap[@↻(L, ux+1), @↺(L, uy-1)]
        hk[Aidx, Bidx] += tprim
        hk[Bidx, Aidx] += adjoint(tprim)
    end; end
    return hk
end


"""
正方双层
"""
function lattice_bilayer_square(::Type{T}, L, hop_t::T, t_prep::T; periodic=true) where T
    Nx = 2*L^2
    hk = zeros(T, Nx, Nx)
    ucmap = zeros(Int, L, L)
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        ucmap[ux, uy] = L*(uy-1) + ux
    end; end
    #
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        #第一层
        Aidx = 2*ucmap[ux, uy]-1
        Bidx = 2*ucmap[@↻(L, ux+1), uy]-1
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #
        Bidx = 2*ucmap[ux, @↻(L, uy+1)]-1
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #第二层
        Aidx = 2*ucmap[ux, uy]
        Bidx = 2*ucmap[@↻(L, ux+1), uy]
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #
        Bidx = 2*ucmap[ux, @↻(L, uy+1)]
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #层间
        Aidx = 2*ucmap[ux, uy]-1
        Bidx = 2*ucmap[ux, uy]
        hk[Aidx, Bidx] = t_prep
        hk[Bidx, Aidx] = adjoint(t_prep)
    end; end
    return hk
end


"""
正方棋盘
"""
function lattice_checkerboard_square(::Type{T}, L, hop_t::T, tprim::T; periodic=true) where T
    Nx = L^2
    hk = zeros(T, Nx, Nx)
    ucmap = zeros(Int, L, L)
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        ucmap[ux, uy] = L*(uy-1) + ux
    end; end
    #
    for ux in Base.OneTo(L); for uy in Base.OneTo(L)
        Aidx = ucmap[ux, uy]
        Bidx = ucmap[@↻(L, ux+1), uy]
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #
        Bidx = ucmap[ux, @↻(L, uy+1)]
        hk[Aidx, Bidx] = hop_t
        hk[Bidx, Aidx] = adjoint(hop_t)
        #
        if mod(ux+uy, 2) == 0
            Bidx = ucmap[@↻(L, ux+1), @↻(L, uy+1)]
            hk[Aidx, Bidx] += tprim
            hk[Bidx, Aidx] += adjoint(tprim)
        else
            Bidx = ucmap[@↻(L, ux+1), @↺(L, uy-1)]
            hk[Aidx, Bidx] += tprim
            hk[Bidx, Aidx] += adjoint(tprim)
        end
    end; end
    return hk
end
